#include "inorder/pipeline.h"
#include "inorder/cache.h"

const char *op_name[NUM_OP] = {"NOP", "ADD", "CBR", "JMP", "MUL", "DIV", "LD", "ST"};

void Pipeline::init_regfile()
{
    for (int i = 0; i < 32; i++)
    {
        regfile[i].valid = true;
    }
}

void Pipeline::init_op_latency()
{
    for (int i = 0; i < NUM_OP; i++)
    {
        op_latency[i] = 1;
        op_cnt[i] = 0;
    }
    op_latency[OP_MUL] = 33;
    op_latency[OP_DIV] = 65;
}

void Pipeline::print_counter()
{
    printf(COLOR("[Pipeline] ", GREEN));
    for (size_t i = 0; i < NUM_OP; i++)
    {
        printf("%s:%ld ", op_name[i], op_cnt[i]);
    }
    printf("\n");
    float ipc = (float)inst_cnt / cycle;
    printf(COLOR("[Pipeline] ", GREEN) "IPC:%.03f\n", ipc);
}

void Pipeline::print_status()
{
    std::cout << "************************\n";

    std::cout << " IFID: ";
    if (IFID.op)
        std::cout << IFID.op->id;
    else
        std::cout << "##";

    std::cout << " IDEX: ";
    if (IDEX.op)
        std::cout << IDEX.op->id;
    else
        std::cout << "##";

    std::cout << " EXMEM: ";
    if (EXMEM.op)
        std::cout << EXMEM.op->id;
    else
        std::cout << "##";

    std::cout << " MEMWB: ";
    if (MEMWB.op)
        std::cout << MEMWB.op->id;
    else
        std::cout << "##";

    std::cout << std::endl
              << "************************\n";
}
Pipeline::Pipeline() : forward_stall("stall_forward"), ex_stall("stall_ex"), mem_stall("stall_mem"), if_stall("stall_if"), cycle("cycle"), inst_cnt("instructions")
{
    init_op_latency();
    init_regfile();
    IFID.valid = false;
    IFID.op = NULL;
    IDEX.valid = false;
    IDEX.op = NULL;
    EXMEM.valid = false;
    EXMEM.op = NULL;
    MEMWB.valid = false;
    MEMWB.op = NULL;
}
Pipeline::~Pipeline()
{
}

void Pipeline::FetchStage()
{
    if (IFID.valid)
        return;

    if (IFID.op && !IFID.valid)
    {
        if_stall++;
        if (IFID.latency > 1)
        {
            IFID.latency--;
        }
        else
        {
            IFID.valid = true;
        }
    }

    if (IFID.op == NULL)
    {
        auto op = read_op();
        IFID.ebreak = !op->valid;
        IFID.op = op;

        if (op->trace.opcode < NUM_OP)
            op_cnt[op->trace.opcode]++;

        bool branch = op->trace.opcode == OP_CBR;
        bool pred_taken = op->trace.branch_pc < op->trace.pc;
        bool branch_miss = branch & (pred_taken != op->trace.branch_taken);
        bool branch_hit = branch & (pred_taken == op->trace.branch_taken);

        if (branch_miss)
        {
            IFID.latency = 2;
            IFID.valid = false;
        }
        else if (branch_hit)
        {
            IFID.latency = 1;
            IFID.valid = false;
        }
        else
        {
            IFID.latency = 0;
            IFID.valid = true;
        }

        branch_miss = op->trace.branch_taken;
        if (branch_miss)
        {
            IFID.latency = 2;
            IFID.valid = false;
        }
        else
        {
            IFID.latency = 0;
            IFID.valid = true;
        }
    }
}

void Pipeline::DecodeStage()
{
    if (IDEX.valid)
        return;
    if (!IFID.valid)
        return;
    IDEX.op = IFID.op;
    IDEX.valid = IFID.valid;
    IDEX.latency = 0;
    IDEX.ebreak = IFID.ebreak;
    IFID.valid = false;
    IFID.op = NULL;
}
void Pipeline::ExecuteStage()
{
    if (EXMEM.valid)
        return;
    if (!IDEX.valid)
        return;

    if (EXMEM.op && !EXMEM.valid)
    {
        /* Wait for DCache return data. */
        if (IDEX.latency > 1)
        {
            IDEX.latency--;
            ex_stall++;
        }
        else
        {
            IDEX.valid = false;
            IDEX.op = NULL;
            EXMEM.latency = 0;
            EXMEM.valid = true;
        }
        return;
    }

    auto op = IDEX.op;
    EXMEM.op = op;
    EXMEM.ebreak = IDEX.ebreak;
    for (int i = 0; i < 2; i++)
    {
        if (op->trace.src[i] > 0)
        {
            if (regfile[op->trace.src[i]].valid == false)
            {
                forward_stall++;
                return;
            }
        }
    }
    int latency = op_latency[op->trace.opcode];
    IDEX.latency = latency - 1;

    if (op->trace.dst > 0)
    {
        /* can't bypass */
        if (latency > 1 || op->trace.opcode == OP_LD)
        {
            regfile[op->trace.dst].valid = false;
        }
    }

    if (IDEX.latency > 0)
    {
        EXMEM.valid = false;
        ex_stall++;
        return;
    }

    EXMEM.latency = 0;
    EXMEM.valid = true;
    IDEX.valid = false;
    IDEX.op = NULL;
}
void Pipeline::MemoryStage()
{
    // assert(MEMWB.valid == false);
    if (MEMWB.valid)
        return;
    if (!EXMEM.valid)
        return;

    /* Wait for DCache access data. */
    // if (MEMWB.op && !MEMWB.valid)
    // {
    //     if (EXMEM.latency > 0)
    //     {
    //         EXMEM.latency--;
    //         mem_stall++;
    //     }
    //     else
    //     {
    //         EXMEM.op = NULL;
    //         EXMEM.valid = false;
    //         MEMWB.latency = 0;
    //         MEMWB.valid = true;
    //     }
    //     return;
    // }

    auto op = EXMEM.op;
    MEMWB.op = op;
    MEMWB.ebreak = EXMEM.ebreak;

    if (op->trace.opcode == OP_LD || op->trace.opcode == OP_ST)
    {
        bool success = cache->send(op);
        if (success)
        {

            EXMEM.latency = INT32_MAX;
            MEMWB.valid = true;
            EXMEM.op = NULL;
            EXMEM.valid = false;
        }
        else
        {
        }
        // bool hit = cache_access(op->trace.mem_access_addr);
        // if (hit)
        // {
        //     EXMEM.latency = 0;
        //     MEMWB.valid = true;
        //     EXMEM.op = NULL;
        //     EXMEM.valid = false;
        // }
        // else
        // {
        //     EXMEM.latency = 7;
        //     MEMWB.valid = false;
        //     mem_stall++;
        // }
    }
    // else if (op->trace.opcode == OP_ST)
    // {
    //     bool hit = cache_access(op->trace.mem_access_addr);
    //     if (hit)
    //     {
    //         EXMEM.latency = 2;
    //     }
    //     else
    //     {
    //         EXMEM.latency = 2 + 7;
    //     }
    //     mem_stall++;
    //     MEMWB.valid = false;
    // }
    else
    {
        MEMWB.latency = 0;
        MEMWB.valid = true;
        EXMEM.op = NULL;
        EXMEM.valid = false;
    }
}
void Pipeline::WritebackStage()
{
    if (!MEMWB.valid)
        return;
    if (MEMWB.op->trace.opcode == OP_LD || MEMWB.op->trace.opcode == OP_ST)
    {
        if (cache->recv(MEMWB.op) == false)
        {
            mem_stall++;
            return;
        }
    }

    auto op = MEMWB.op;

    static int8_t pending_dst = -1;
    if (pending_dst >= 0)
        regfile[pending_dst].valid = true;
    pending_dst = op->trace.dst;
    if (MEMWB.ebreak)
        sim_finish = true;

    inst_cnt++;
    // if (op->trace.opcode < NUM_OP)
    //     op_cnt[op->trace.opcode]++;
    free_op(op);
    MEMWB.op = NULL;
    MEMWB.valid = false;
}

void Pipeline::run()
{

    WritebackStage();
    MemoryStage();
    ExecuteStage();
    DecodeStage();
    FetchStage();

    // print_status();
    cycle++;
}